package com.bytecode;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * <p>
 * 字节码转译
 * </p>
 *
 * @author Arnold Yang
 * @since 2018-09-03
 */
public class Transformer implements ClassFileTransformer {

  protected Logger logger = LoggerFactory.getLogger(Transformer.class);

  // 不需要处理的class
  private static final Pattern CLASS_SKIP_PATTERN = Pattern.compile("^(sun/)|(com/sun)|(org/)|(junit/)|(com/intellij).*$");
  private final Matcher classMatcher = CLASS_SKIP_PATTERN.matcher("");

  // 不需要处理的method
  private static final Pattern METHOD_SKIP_PATTERN = Pattern.compile("^(java\\.).*$");
  private final Matcher methodMatcher = METHOD_SKIP_PATTERN.matcher("");

  @Override
  public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws RuntimeException {
    if(loader == null | className == null){
      return null;
    }

    if(logger.isDebugEnabled()){
      //logger.debug("load class:{}", className);
    }

    if(skipClass(className)){
      return null;
    }

    ClassPool pool = new ClassPool(true);
    pool.appendClassPath(new LoaderClassPath(loader));
    try {
      CtClass cls = pool.makeClass(new ByteArrayInputStream(classfileBuffer));
      //CtClass cls = pool.getCtClass(className);

      CtMethod[] methods = cls.getMethods();
      if(methods != null && methods.length > 0){
        for(CtMethod method : methods){
          if(skipMethod(method.getLongName())){
            continue;
          }
          method.addLocalVariable("startTime", CtClass.longType);
          method.insertBefore("startTime = System.currentTimeMillis();");
          method.insertAfter(String.format("System.out.println(\"call method %s cost(ms):\"+(System.currentTimeMillis() - startTime));", method.getLongName()));
        }
      }

      return cls.toBytecode();
    } catch (Exception e) {
      logger.error("", e);
    }

    return new byte[0];
  }

  private boolean skipMethod(String methodName) {
    methodMatcher.reset(methodName);
    return methodMatcher.find();
  }

  private boolean skipClass(String className) {
    classMatcher.reset(className);
    return classMatcher.find();
  }
}
